#!/usr/bin/env python3
"""
SD3采样参数节点
专门管理SD3训练中的采样相关配置
只包含SD3真正支持的采样参数
"""

import os
from .lora_trainer_utils.constants import SD3_TRAINING_PARAMS_CATEGORY

class SD3SamplingParams:
    """SD3采样参数节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # 基础采样设置
                "sample_every_n_steps": ("INT", {
                    "default": 200,
                    "min": 0,
                    "max": 10000,
                    "tooltip": "每多少步采样输出一次样本图像，0为不采样"
                }),
                "sample_every_n_epochs": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 100,
                    "tooltip": "每多少epoch采样输出一次样本图像，0为不采样"
                }),
                "sample_prompts": ("STRING", {
                    "default": "",
                    "tooltip": "采样用prompt文件路径（.txt格式）"
                }),
                "sample_sampler": (["ddim", "pndm", "lms", "euler", "euler_a", "heun", "dpm_2", "dpm_2_a", "dpmsolver", "dpmsolver++", "dpmsingle", "k_lms", "k_euler", "k_euler_a", "k_dpm_2", "k_dpm_2_a"], {
                    "default": "ddim",
                    "tooltip": "采样器算法"
                }),
            },
            "optional": {
                "sample_at_first": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "训练前生成样本图像"
                }),
            }
        }
    
    @classmethod
    def VALIDATE_INPUTS(cls, **kwargs):
        """验证输入参数"""
        # 检查采样提示词文件（如果提供）
        sample_prompts = kwargs.get("sample_prompts")
        if sample_prompts and isinstance(sample_prompts, str) and sample_prompts.strip():
            if not os.path.exists(sample_prompts):
                return f"采样提示词文件不存在: {sample_prompts}"
        
        return True
    
    RETURN_TYPES = ("SD3_SAMPLING_PARAMS",)
    RETURN_NAMES = ("采样参数",)
    FUNCTION = "generate_sampling_params"
    CATEGORY = SD3_TRAINING_PARAMS_CATEGORY
    
    def generate_sampling_params(self, **kwargs):
        """生成采样参数"""
        
        # 智能参数调整
        sampling_params = {}
        
        # 处理采样间隔
        sample_every_n_steps = kwargs.get("sample_every_n_steps", 0)
        sample_every_n_epochs = kwargs.get("sample_every_n_epochs", 0)
        
        if sample_every_n_steps == 0 and sample_every_n_epochs == 0:
            print("[WARNING] 采样间隔设置为0，将禁用采样功能")
        else:
            print(f"[INFO] 采样间隔: 每{sample_every_n_steps}步或每{sample_every_n_epochs}轮")
        
        # 复制所有参数
        for key, value in kwargs.items():
            sampling_params[key] = value
        
        # 采样配置建议
        print("[INFO] 📸 SD3采样参数配置:")
        if sampling_params.get("sample_every_n_steps", 0) > 0:
            print(f"   ✅ 已启用采样，间隔: {sampling_params['sample_every_n_steps']}步")
        elif sampling_params.get("sample_every_n_epochs", 0) > 0:
            print(f"   ✅ 已启用采样，间隔: 每{sampling_params['sample_every_n_epochs']}轮")
        else:
            print("   ⚠️ 未启用采样功能")
        
        if sampling_params.get("sample_prompts"):
            print(f"   📝 使用采样提示词文件: {sampling_params['sample_prompts']}")
        else:
            print("   💡 建议提供采样提示词文件以生成有意义的样本")
        
        print(f"   🎲 采样器算法: {sampling_params['sample_sampler']}")
        
        if sampling_params.get("sample_at_first", False):
            print("   ✅ 已启用训练前采样")
        
        return (sampling_params,)


# 注册节点
NODE_CLASS_MAPPINGS = {
    "SD3SamplingParams": SD3SamplingParams,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SD3SamplingParams": "📸 SD3采样参数",
} 